package cn.remex.configuration.etcd.client;

import cn.remex.configuration.etcd.schema.*;
import cn.remex.core.util.Judgment;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.protobuf.ByteString;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.StatusRuntimeException;

import javax.annotation.Nullable;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
 * Disc:
 * Created by xinxinran on 2017/7/18.
 */
public class EtcdV3Client {

    private final int maxWatchCount;
    private final ExecutorService watchPool;

    private final ManagedChannel channel;
    private final KVGrpc.KVFutureStub kvStub;
    private final WatchGrpc.WatchBlockingStub watchStub;

    public EtcdV3Client(String connString){
        String[] connInfo = connString.split(":");
        String host = connInfo[0];
        int port = 2379;
        int maxWatchCount = 20;
        if(connInfo.length > 1 && connInfo[1].matches("^\\d{2,6}$")){
            port = Integer.parseInt(connInfo[1]);
        }

        if(connInfo.length > 2 && connInfo[2].matches("^\\d{2,6}$")){
            maxWatchCount = Integer.parseInt(connInfo[2]);
        }

        this.channel = ManagedChannelBuilder.forAddress(host, port)
                .usePlaintext(true).build();
        kvStub = KVGrpc.newFutureStub(channel);
        watchStub = WatchGrpc.newBlockingStub(channel);

        this.maxWatchCount = maxWatchCount;
        if(maxWatchCount > 0){
            watchPool = Executors.newFixedThreadPool(maxWatchCount);
        }else{
            watchPool = null;
        }

    }

    public EtcdV3Client(String host, int port, int maxWatchCount){
        this.channel = ManagedChannelBuilder.forAddress(host, port)
                .usePlaintext(true).build();
        kvStub = KVGrpc.newFutureStub(channel);
        watchStub = WatchGrpc.newBlockingStub(channel);

        this.maxWatchCount = maxWatchCount;
        if(maxWatchCount > 0){
            watchPool = Executors.newFixedThreadPool(maxWatchCount);
        }else{
            watchPool = null;
        }
    }

    public CompletableFuture<PutResponse> setKV(String key, String value){

        if(Judgment.nullOrBlank(key)){
            throw new EtcdV3ClientException("[ETCD::setKV] key can not be null or empty");
        }

        if(value == null){
            value = "";
        }

        ListenableFuture<PutResponse> future =  kvStub.put(PutRequest.newBuilder()
                .setKey(ByteString.copyFromUtf8(key))
                .setValue(ByteString.copyFromUtf8(value))
                .build());

        CompletableFuture<PutResponse> resultFuture = new CompletableFuture<>();
        Futures.addCallback(future, new FutureCallback<PutResponse>() {
            @Override
            public void onSuccess(@Nullable PutResponse putResponse) {
                resultFuture.complete(putResponse);
            }

            @Override
            public void onFailure(Throwable throwable) {
                resultFuture.completeExceptionally(throwable);
            }
        });

        return resultFuture;
    }

    public CompletableFuture<String> getValue(String key){
        if(Judgment.nullOrBlank(key)){
            throw new EtcdV3ClientException("[ETCD::getValue] key can not be null or empty");
        }

        ListenableFuture<RangeResponse> future = kvStub.range(RangeRequest.newBuilder()
                .setKey(ByteString.copyFromUtf8(key))
                .build());

        CompletableFuture<String> resultFuture = new CompletableFuture<>();
        Futures.addCallback(future, new FutureCallback<RangeResponse>() {
            @Override
            public void onSuccess(@Nullable RangeResponse rangeResponse) {
                if(rangeResponse == null || rangeResponse.getCount() == 0){
                    resultFuture.complete(null);
                }else{
                    resultFuture.complete(rangeResponse.getKvs(0).getValue().toStringUtf8());
                }
            }

            @Override
            public void onFailure(Throwable throwable) {
                resultFuture.completeExceptionally(throwable);
            }
        });

        return resultFuture;
    }

    public CompletableFuture<Map<String, String>> getRange(String keyPrefix){
        if(Judgment.nullOrBlank(keyPrefix) || "/".equals(keyPrefix)){
            throw new EtcdV3ClientException("[ETCD::getRange] key can not be null or empty");
        }

        if(keyPrefix.endsWith("/")){
            keyPrefix = keyPrefix.substring(0, keyPrefix.length() - 1);
        }
        ListenableFuture<RangeResponse> future = kvStub.range(RangeRequest.newBuilder()
                .setKey(ByteString.copyFromUtf8(keyPrefix))
                .setRangeEnd(ByteString.copyFromUtf8(getNextKeyForRange(keyPrefix)))
                .build());

        CompletableFuture<Map<String, String>> resultFuture = new CompletableFuture<>();
        Futures.addCallback(future, new FutureCallback<RangeResponse>() {
            @Override
            public void onSuccess(@Nullable RangeResponse rangeResponse) {
                if(rangeResponse == null || rangeResponse.getCount() == 0){
                    resultFuture.complete(new HashMap<>());
                }else{
                    Map<String, String> resultMap =  rangeResponse.getKvsList().stream().collect(Collectors.toMap(keyValue -> keyValue.getKey().toStringUtf8(), keyValue -> keyValue.getValue().toStringUtf8()));
                    resultFuture.complete(resultMap);
                }
            }

            @Override
            public void onFailure(Throwable throwable) {
                resultFuture.completeExceptionally(throwable);
            }
        });
        return resultFuture;
    }

    private String getNextKeyForRange(String keyPrefix){
        byte[] keyBytes = keyPrefix.getBytes();

        keyBytes[keyBytes.length - 1] = (byte)(keyBytes[keyBytes.length - 1] + 1);

        return new String(keyBytes);
    }

    private int validNoBlockingWatchCount(){
        if(this.watchPool == null){
            return 0;
        }

        return this.maxWatchCount - ((ThreadPoolExecutor)this.watchPool).getActiveCount();
    }

    public void watchOne(String key, Consumer<WatchResponse> processor, boolean needBlocked){
        watch(key, null, processor, needBlocked);
    }

    public void watchRange(String keyPrefix, Consumer<WatchResponse> processor, boolean needBlocked){
        watch(keyPrefix, getNextKeyForRange(keyPrefix), processor, needBlocked);
    }

    private void watch(String key, String rangeEnd, Consumer<WatchResponse> processor, boolean needBlocked){
        WatchCreateRequest.Builder wcrBuilder = WatchCreateRequest.newBuilder()
                .setKey(ByteString.copyFromUtf8(key));
        if(!Judgment.nullOrBlank(rangeEnd)){
            wcrBuilder.setRangeEnd(ByteString.copyFromUtf8(rangeEnd));
        }
        Iterator<WatchResponse> iterator = watchStub.watch(WatchRequest.newBuilder()
                .setCreateRequest(wcrBuilder)
                .build());

        if(needBlocked){
            try {
                while (iterator.hasNext()) {
                    processor.accept(iterator.next());
                }
            }catch(StatusRuntimeException e){
                boolean sleepFinished = false;
                try {
                    Thread.sleep(5000);
                    sleepFinished = true;
                } catch (InterruptedException ignored) { sleepFinished = false; }
                if (sleepFinished) watch(key, rangeEnd, processor, true);
            }
        }else{
            if(validNoBlockingWatchCount() > 0){
                this.watchPool.submit(() -> {
                    try {
                        while(iterator.hasNext()){
                            processor.accept(iterator.next());
                        }
                    }catch(StatusRuntimeException e){
                        boolean sleepFinished = false;
                        try {
                            Thread.sleep(15000);
                            sleepFinished = true;
                        } catch (InterruptedException ignored) { sleepFinished = false; }
                        if (sleepFinished) watch(key, rangeEnd, processor, true);
                    }
                });
            }else{
                throw new EtcdV3ClientException("[ETCD] no valid thread for nonblocking watching...");
            }
        }

    }


    public void shutdown() throws InterruptedException {
        channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
    }

    public static void main(String[] args) throws Exception {
        EtcdV3Client client = new EtcdV3Client("127.0.0.1", 2379, 10);

//        client.setKV("/testz/a", "aaa").thenComposeAsync(putResponse -> {
//            System.out.println(putResponse);
//            return client.setKV("/testz/b", "bbb");
//        }).thenComposeAsync(putResponse -> {
//            System.out.println(putResponse);
//            return client.getRange("/testz");
//        }).thenAccept(System.out::println).exceptionally(throwable -> {
//            throwable.printStackTrace();
//            return null;
//        }).get();

//        client.watch("/test/a", System.out::println, false);
        client.watchRange("/test", System.out::println, true);

        client.shutdown();
    }
}
